import torch
import paras
import LSTM
import LSTMModule
from matplotlib import gridspec
from mpl_finance import candlestick_ochl
from pylab import *
import autocode_torch
from madin import joinKline
from mongoapi import *

device = torch.device('cpu')


def get_EMA(df, a=0.7):
    ema = df
    for i in range(len(df)):
        if i == 0:
            ema[i] = df[i]
        if i > 0:
            ema[i] = (1 - a) * ema[i - 1] + a * df[i]
    return ema


def Show(datalist, svcdatas, span=200):
    svcdata = svcdatas[0]
    spandata = datalist[-span:]

    datas = []
    for i in range(len(spandata)):
        spandata[i]["pos"] = i
        datas.append([i, spandata[i]["Open"], spandata[i]["Close"], spandata[i]["High"], spandata[i]["Low"],
                      spandata[i]["Date"]])
    idtodate = {}
    idtodate = {one[0]: one[5] for one in datas}
    showid = [one[0] for one in datas]
    thred = 0.99 * max(svcdata)
    thred2 = 1.01 * min(svcdata)
    point0 = map(
            lambda i: spandata[i]["Low"] * (1 - 3e-4) if (svcdatas[1, i] <= thred2 and i > 0) else None,
            range(len(svcdata)))
    point1 = map(
            lambda i: spandata[i]["High"] * (1 + 3e-4) if (svcdatas[2, i] >= thred and i > 0) else None,
            range(len(svcdata)))
    # point00 = map(lambda i: spandata[i]["Low"] + 0.1 if svcdata1[i] > 0.05 > svcdata1[i - 1] else None,
    #               range(len(svcdata)))
    # point11 = map(lambda i: spandata[i]["Low"] - 0.1 if svcdata1[i] < -0.05 < svcdata1[i - 1] else None,
    #               range(len(svcdata)))

    fig = plt.figure()
    fig.subplots_adjust(top=0.98, bottom=0.02, left=0.05, right=0.98)

    fig.subplots_adjust(hspace=0.01)
    gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1])
    ax = plt.subplot(gs[0])
    ax.set_ylabel("USD/oz")
    ax1 = plt.subplot(gs[1], sharex=ax)

    def getDate(x, position):
        if x in idtodate:
            return idtodate[x]
        return ""

    ax.xaxis.set_major_formatter(FuncFormatter(getDate))
    candlestick_ochl(ax, datas, width=0.5, colorup='r', colordown='g')
    ax.plot(showid, list(point0), color="r", marker='o')
    ax.plot(showid, list(point1), color="g", marker='o')
    # ax.plot(showid, list(point00), color="r", marker='o')
    # ax.plot(showid, list(point11), color="g", marker='o')

    # ax1.plot(showid, svcdata2-svcdata1, color="y")
    color = ['r', 'b', 'y', 'm', 'g', 'c']
    for i in range(1, len(svcdatas)):
        ax1.plot(showid, svcdatas[i], color=color[i])
    ax1.hlines([0, thred], showid[0], showid[-1], colors=['b', 'g', 'g'], linestyles='dashdot')

    ax.autoscale_view()
    ax1.autoscale_view()

    ax.grid()
    ax1.grid()

    mng = plt.get_current_fig_manager()
    # mng.full_screen_toggle()
    plt.show()


def getrnn(datalists, span=100, fit=1):
    disres = []
    rnn = torch.load('data/{}_{}_lstm.pkl'.format(paras.symbol, paras.qt_type))
    for i in range(100 + LSTMModule.TIME_STEP, len(datalists)):
        xtest = datalists[i - 100 - LSTMModule.TIME_STEP:i]
        datacalc = [LSTMModule.getone(xtest[:i + 1]) for i in range(100, len(xtest))]
        testdata = torch.from_numpy(np.array(datacalc, dtype=np.float32)).view(-1, LSTMModule.TIME_STEP,
                                                                               LSTMModule.INPUT_SIZE)
        disre = rnn(testdata)
        disre = disre.data.numpy()
        disres.append(disre[-1])
    ar = np.array(disres)
    ar = torch.max(torch.from_numpy(ar), 1)[1].cpu().data.numpy()
    return ar


def showRnn(dbname="FX", symbol='FXUSDJPY', qt_type=5):
    datalist = finddb(dbname=dbname, tablename=symbol, limit=200000)
    # datalist = json.load(open("data/A%s.json" % symbol, 'r'))
    datalist = joinKline(datalist, qt_type)
    datalist = datalist[-60000:]

    print(datetime.datetime.now())
    disres = getrnn(datalist)
    print(datetime.datetime.now())
    ds = [disres, disres, disres]
    ll = min([len(o) for o in ds])
    ds = np.array([o[-ll:] for o in ds])
    datalist = datalist[-ll:]
    print(ds.shape)
    Show(datalist, ds, span=len(ds[0]))


if __name__ == "__main__":
    # autocode_torch.autocode(paras.dbname, paras.symbol, paras.qt_type)
    LSTM.updatemodel(paras.dbname, paras.symbol, paras.qt_type)
    showRnn(paras.dbname, paras.symbol, paras.qt_type)
